#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket 处理器
"""

import json
import asyncio
from typing import Dict
from fastapi import WebSocket, WebSocketDisconnect

# 导入认证相关模块
import sys
from pathlib import Path
current_dir = Path(__file__).parent
parent_dir = current_dir.parent.parent
sys.path.insert(0, str(parent_dir))

try:
    from security_manager import get_password, store_password
except ImportError as e:
    print(f"导入安全管理器失败: {e}")

class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket, client_id: str):
        """建立WebSocket连接"""
        await websocket.accept()
        self.active_connections[client_id] = websocket

    def disconnect(self, client_id: str):
        """断开WebSocket连接"""
        if client_id in self.active_connections:
            del self.active_connections[client_id]

    async def send_message(self, client_id: str, message: dict):
        """发送消息到指定客户端"""
        if client_id in self.active_connections:
            try:
                await self.active_connections[client_id].send_text(json.dumps(message))
            except:
                self.disconnect(client_id)

class WebSocketHandler:
    """WebSocket消息处理器"""
    
    def __init__(self, manager: ConnectionManager):
        self.manager = manager
    
    async def handle_message(self, client_id: str, message: dict):
        """处理WebSocket消息"""
        try:
            if 'action' in message:
                await self._handle_action(client_id, message)
            elif message.get('type'):
                await self._handle_type_message(client_id, message)
            else:
                await self.manager.send_message(client_id, {"type": "ping", "message": "pong"})
        except Exception as e:
            await self.manager.send_message(client_id, {
                "type": "error", 
                "message": f"处理消息时发生错误: {str(e)}"
            })
    
    async def _handle_action(self, client_id: str, message: dict):
        """处理带action的消息"""
        action = message['action']
        
        if action == 'verify_premium_password':
            await self._verify_premium_password(client_id, message)
        elif action == 'change_premium_password':
            await self._change_premium_password(client_id, message)
    
    async def _handle_type_message(self, client_id: str, message: dict):
        """处理带type的消息"""
        msg_type = message['type']
        
        if msg_type == 'apply_cookies':
            await self._apply_cookies(client_id, message)
        elif msg_type == 'apply_general_settings':
            await self._apply_general_settings(client_id, message)
        elif msg_type == 'apply_proxy_settings':
            await self._apply_proxy_settings(client_id, message)
        elif msg_type == 'test_proxy_connection':
            await self._test_proxy_connection(client_id, message)
        elif msg_type == 'change_premium_password':
            await self._change_premium_password(client_id, message)
    
    async def _verify_premium_password(self, client_id: str, message: dict):
        """验证付费区密码"""
        password = message.get('password', '')
        stored_password = get_password("premium_access") or "520184"
        
        success = password == stored_password
        
        await self.manager.send_message(client_id, {
            "type": "premium_password_verification",
            "success": success,
            "message": "密码验证成功" if success else "密码验证失败"
        })
    
    async def _change_premium_password(self, client_id: str, message: dict):
        """修改付费区密码"""
        current_password = message.get('current_password', '')
        new_password = message.get('new_password', '')
        
        # 验证输入
        if not current_password or not new_password:
            await self.manager.send_message(client_id, {
                "type": "premium_password_changed",
                "success": False,
                "message": "当前密码和新密码不能为空"
            })
            return
        
        if len(new_password) < 4:
            await self.manager.send_message(client_id, {
                "type": "premium_password_changed",
                "success": False,
                "message": "新密码长度至少为4位"
            })
            return
        
        if current_password == new_password:
            await self.manager.send_message(client_id, {
                "type": "premium_password_changed",
                "success": False,
                "message": "新密码不能与当前密码相同"
            })
            return
        
        stored_password = get_password("premium_access") or "520184"
        
        if current_password == stored_password:
            # 密码正确，可以修改
            if store_password("premium_access", new_password):
                await self.manager.send_message(client_id, {
                    "type": "premium_password_changed",
                    "success": True,
                    "message": "密码修改成功"
                })
            else:
                await self.manager.send_message(client_id, {
                    "type": "premium_password_changed",
                    "success": False,
                    "message": "密码存储失败"
                })
        else:
            # 当前密码错误
            await self.manager.send_message(client_id, {
                "type": "premium_password_changed",
                "success": False,
                "message": "当前密码错误"
            })
    
    async def _apply_cookies(self, client_id: str, message: dict):
        """应用Cookie设置"""
        # 这里可以添加实际的Cookie应用逻辑
        await self.manager.send_message(client_id, {
            "type": "cookies_applied",
            "success": True,
            "message": "Cookie设置已应用"
        })
    
    async def _apply_general_settings(self, client_id: str, message: dict):
        """应用通用设置"""
        # 这里可以添加实际的设置应用逻辑
        await self.manager.send_message(client_id, {
            "type": "general_settings_applied",
            "success": True,
            "message": "通用设置已应用"
        })
    
    async def _apply_proxy_settings(self, client_id: str, message: dict):
        """应用代理设置"""
        # 这里可以添加实际的代理设置逻辑
        await self.manager.send_message(client_id, {
            "type": "proxy_settings_applied",
            "success": True,
            "message": "代理设置已应用"
        })
    
    async def _test_proxy_connection(self, client_id: str, message: dict):
        """测试代理连接"""
        # 这里可以添加实际的代理测试逻辑
        await self.manager.send_message(client_id, {
            "type": "proxy_test_result",
            "success": True,
            "message": "代理连接测试成功"
        })

# 全局管理器实例
manager = ConnectionManager()
websocket_handler = WebSocketHandler(manager)

async def websocket_endpoint(websocket: WebSocket, client_id: str):
    """WebSocket端点处理函数"""
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_text()
            try:
                message = json.loads(data)
                await websocket_handler.handle_message(client_id, message)
            except json.JSONDecodeError:
                await manager.send_message(client_id, {"type": "ping", "message": "pong"})
    except WebSocketDisconnect:
        manager.disconnect(client_id)
